import torch
from models.temporal_yolo import create_temporal_yolo
import os
from ultralytics import YOLO
from torch.utils.data import Dataset, DataLoader
import cv2
import numpy as np
from pathlib import Path
from torch.optim import Adam
from tqdm import tqdm

class TemporalDataset(Dataset):
    def __init__(self, data_dir, sequence_length=3):
        self.data_dir = Path(data_dir)
        self.sequence_length = sequence_length
        self.image_files = sorted(list(self.data_dir.glob('*.jpg')))
        self.sequences = self._create_sequences()
        
    def _create_sequences(self):
        """创建连续帧序列"""
        sequences = []
        for i in range(len(self.image_files) - self.sequence_length + 1):
            seq = self.image_files[i:i + self.sequence_length]
            sequences.append(seq)
        return sequences
    
    def __len__(self):
        return len(self.sequences)
    
    def __getitem__(self, idx):
        seq_paths = self.sequences[idx]
        frames = []
        labels = []
        
        for path in seq_paths:
            # 读取图像
            img = cv2.imread(str(path))
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = cv2.resize(img, (640, 640))  # 调整大小
            img = img.transpose(2, 0, 1)  # HWC to CHW
            img = img / 255.0  # 归一化
            
            # 读取对应的标签
            label_path = path.with_suffix('.txt')
            if label_path.exists():
                with open(label_path, 'r') as f:
                    label = np.array([list(map(float, line.strip().split())) 
                                    for line in f.readlines()])
            else:
                label = np.array([])
            
            frames.append(img)
            labels.append(label)
        
        # 转换为numpy数组
        frames = np.stack(frames)
        
        return torch.from_numpy(frames).float(), labels

def train_one_epoch(model, train_loader, optimizer, device, epoch):
    model.train()
    total_loss = 0
    progress_bar = tqdm(train_loader, desc=f'Epoch {epoch}')
    
    for batch_idx, (frames, targets) in enumerate(progress_bar):
        frames = frames.to(device)
        
        # 前向传播
        loss = model.train_step((frames, targets), device)
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
        
        # 更新进度条
        progress_bar.set_postfix({'loss': loss.item()})
    
    return total_loss / len(train_loader)

def main():
    # 获取当前目录的绝对路径
    current_dir = os.path.dirname(os.path.abspath(__file__))
    
    # 创建数据集和数据加载器
    train_dataset = TemporalDataset(os.path.join(current_dir, 'processed_dataset/train/images'))
    val_dataset = TemporalDataset(os.path.join(current_dir, 'processed_dataset/val/images'))
    
    train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_dataset, batch_size=4, num_workers=4)
    
    # 创建时序增强的YOLO模型
    model = create_temporal_yolo()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    
    # 优化器
    optimizer = Adam(model.parameters(), lr=0.001)
    
    # 训练循环
    num_epochs = 150
    best_loss = float('inf')
    
    for epoch in range(num_epochs):
        # 训练一个epoch
        train_loss = train_one_epoch(model, train_loader, optimizer, device, epoch)
        print(f'Epoch {epoch}, Average Loss: {train_loss:.4f}')
        
        # 保存检查点
        if train_loss < best_loss:
            best_loss = train_loss
            checkpoint = {
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'loss': train_loss,
            }
            torch.save(checkpoint, f'pickleball_pose_temporal/best_model.pt')
        
        # 每10个epoch保存一次
        if (epoch + 1) % 10 == 0:
            checkpoint = {
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'loss': train_loss,
            }
            torch.save(checkpoint, f'pickleball_pose_temporal/checkpoint_epoch_{epoch+1}.pt')

if __name__ == "__main__":
    main() 